package main

import (
	"fmt"
	"github.com/gorilla/mux"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		// 在这里进行源站验证, 例如检查 Origin 头
		return true
	},
}

func main() {
	r := mux.NewRouter()
	r.HandleFunc("/ws/{sid}", handleWebSocket).Methods("GET")
	r.Handle("/", http.FileServer(http.Dir(".")))

	fmt.Println("Server listening on :7070")
	log.Fatal(http.ListenAndServe(":7070", r))
}

func handleWebSocket(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	sessionID := vars["sid"]
	log.Printf("Received WebSocket connection for session ID: %s", sessionID)

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("Failed to upgrade connection:", err)
		return
	}
	defer conn.Close()

	for {
		_, message, err := conn.ReadMessage()
		if err != nil {
			log.Println("Error reading message:", err)
			break
		}
		log.Printf("Received message for session %s: %s", sessionID, message)

		err = conn.WriteMessage(websocket.TextMessage, []byte("Hello from server!"))
		if err != nil {
			log.Println("Error writing message:", err)
			break
		}
	}
}
